﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Configuration;
using System.Diagnostics;
using System.Xml;
using System.IO;
using System.ServiceModel.Channels;
using System.Runtime.Serialization.Json;
using System.ServiceModel;
using System.Reflection;
using System.Threading;
using Spike.ServiceFramework.Common;

namespace Spike.ServiceFramework.ServiceModel.Audit
{
    public class AuditBehavior : ConfigurationElement
    {
        protected const string HEADER_NAMESPACE = "urn:spike.serviceframework.servicemodel:20120524";
        protected const string HEADER_CORRELATION = "Correlation";
        protected const string HEADER_EVENT = "Event";

        private TraceSource _ts;

        private TraceSource Trace
        {
            get { return _ts ?? (_ts = new TraceSource(TraceSource)); }
        }

        [ConfigurationProperty("includePayloads", DefaultValue = false, IsRequired = false)]
        public bool ShouldTraceMessages
        {
            get { return (bool) base["includePayloads"]; }
            set { base["includePayloads"] = value; }
        }

        [ConfigurationProperty("includeCorrelation", DefaultValue = false, IsRequired = false)]
        public bool ShouldCorrelate
        {
            get { return (bool)base["includeCorrelation"]; }
            set { base["includeCorrelation"] = value; }
        }

        [ConfigurationProperty("propagateCorrelation", DefaultValue = false, IsRequired = false)]
        public bool ShouldPropagateCorrelation
        {
            get { return (bool)base["propagateCorrelation"]; }
            set { base["propagateCorrelation"] = value; }
        }

        [ConfigurationProperty("includeEvent", DefaultValue = false, IsRequired = false)]
        public bool ShouldEvent
        {
            get { return (bool)base["includeEvent"]; }
            set { base["includeEvent"] = value; }
        }

        [ConfigurationProperty("propagateEvent", DefaultValue = false, IsRequired = false)]
        public bool ShouldPropagateEvent
        {
            get { return (bool)base["propagateEvent"]; }
            set { base["propagateEvent"] = value; }
        }

        [ConfigurationProperty("traceSource", DefaultValue = "Spike.ServiceFramework.ServiceModel.Audit", IsRequired = false)]
        public string TraceSource
        {
            get { return (string) base["traceSource"]; }
            set { base["traceSource"] = value; }
        }

        #region Service Tracing
        internal ServiceTraceRecord CreateServiceTraceRecord(ref System.ServiceModel.Channels.Message request)
        {
            return new ServiceTraceRecord
                       {                           
                           AfterReceiveRequest = DateTime.Now,
                           ContractName = OperationContext.Current.EndpointDispatcher.ContractName,
                           ContractNamespace = OperationContext.Current.EndpointDispatcher.ContractNamespace,
                           URI = OperationContext.Current.EndpointDispatcher.EndpointAddress.Uri.AbsoluteUri,
                           Request = ShouldTraceMessages ? MessageToString(ref request) : null,
                           ApplicationDomain = Thread.GetDomain().FriendlyName,
                           Directory = Environment.CurrentDirectory,
                           DomainName = Environment.UserDomainName,
                           MachineName = Environment.MachineName,
                           UserName = Environment.UserName,
                           Headers = TraceHeaders(request.Headers),
                           CorrelationID = GetHeaderCorrelationId(request),
                           EventName = GetHeaderEvent(request)
                       };
        }

        internal void CompleteServiceTraceRecord(ref Message reply, object correlationState)
        {
            if (!reply.IsEmpty && !reply.IsFault)
            {
                var trace = correlationState as ServiceTraceRecord;

                if (trace == null)
                {
                    Trace.TraceEvent(TraceEventType.Error, 101, "DispatchMessageInspector received an empty correlation state when auditing was enabled.");
                }
                else
                {
                    trace.BeforeSendReply = DateTime.Now;
                    trace.Reply = ShouldTraceMessages ? MessageToString(ref reply) : null;

                    Trace.TraceData(TraceEventType.Information, 100, trace);
                }
            }
        }
        #endregion

        #region Client Tracing

        internal ClientTraceRecord CreateClientTraceRecord(ref Message request, System.ServiceModel.IClientChannel channel)
        {
            return new ClientTraceRecord
                       {
                           BeforeSendRequest = DateTime.Now,
                           Request = ShouldTraceMessages ? MessageToString(ref request) : null,
                           URI = channel.RemoteAddress.Uri.AbsoluteUri,
                           ApplicationDomain = Thread.GetDomain().FriendlyName,
                           Directory = Environment.CurrentDirectory,
                           DomainName = Environment.UserDomainName,
                           MachineName = Environment.MachineName,
                           UserName = Environment.UserName,
                           Headers = TraceHeaders(request.Headers),
                           CorrelationID = GetHeaderCorrelationId(request),
                           EventName = GetHeaderEvent(request)
                       };
        }

        internal void CompleteClientTraceRecord(ref System.ServiceModel.Channels.Message reply, object correlationState)
        {
            var trace = correlationState as ClientTraceRecord;

            if (trace == null)
            {
                Trace.TraceEvent(TraceEventType.Error, 101, "ClientMessageInspector received an empty correlation state when auditing was enabled.");
            }
            else
            {
                trace.AfterReceiveReply = DateTime.Now;
                trace.Reply = ShouldTraceMessages ? MessageToString(ref reply) : null;

                Trace.TraceData(TraceEventType.Information, 100, trace);
            }
        }

        #endregion

        #region Header Methods
        private static MessageHeader CreateMessageHeader(string name, object content)
        {
            return MessageHeader.CreateHeader(name, HEADER_NAMESPACE, content);
        }

        private static List<string> TraceHeaders(MessageHeaders messageHeaders)
        {
            List<string> headers = new List<string>();
            messageHeaders.ForEach(l => headers.Add(l.ToString()));
            return headers;

        }

        private string GetHeaderEvent(Message request)
        {
            string eventName = string.Empty;
            if (ShouldEvent)
            {
                int headerID = request.Headers.FindHeader(HEADER_EVENT, HEADER_NAMESPACE);

                if (headerID == -1)
                {
                    eventName = GetEventName();
                    request.Headers.Add(CreateMessageHeader(HEADER_EVENT, eventName));
                }
                else
                {
                    eventName = request.Headers.GetHeader<string>(headerID);

                    if (ShouldPropagateEvent)
                        SetEventName(eventName);
                }
            }
            return eventName;
        }

        private Guid GetHeaderCorrelationId(Message request)
        {
            //Guid correlationID = Guid.Empty;
            //if (ShouldCorrelate)
            //{
            //    int headerID = request.Headers.FindHeader(HEADER_CORRELATION, HEADER_NAMESPACE);

            //    if (headerID != -1)
            //    {
            //        correlationID = request.Headers.GetHeader<Guid>(headerID);

            //        if (ShouldPropagateCorrelation)
            //            SetCorrelationID(correlationID);
            //    }
            //}
            Guid correlationID = Guid.Empty;
            if (ShouldCorrelate)
            {
                int headerId = request.Headers.FindHeader(HEADER_CORRELATION, HEADER_NAMESPACE);
                if (headerId == -1)
                {
                    correlationID = GetCorrelationID();

                    request.Headers.Add(CreateMessageHeader(HEADER_CORRELATION, correlationID));
                }
                else
                {
                    correlationID = request.Headers.GetHeader<Guid>(headerId);

                    if (ShouldPropagateCorrelation)
                        SetCorrelationID(correlationID);
                }
            }
            return correlationID;
        }
        #endregion

        #region http://code.msdn.microsoft.com/windowsdesktop/WCF-REST-Message-Inspector-c4b6790b
        private WebContentFormat GetMessageContentFormat(Message message)
        {
            WebContentFormat format = WebContentFormat.Default;
            if (message.Properties.ContainsKey(WebBodyFormatMessageProperty.Name))
            {
                WebBodyFormatMessageProperty bodyFormat;
                bodyFormat = (WebBodyFormatMessageProperty) message.Properties[WebBodyFormatMessageProperty.Name];
                format = bodyFormat.Format;
            }

            return format;
        }

        protected string MessageToString(ref Message message)
        {
            WebContentFormat messageFormat = this.GetMessageContentFormat(message);
            MemoryStream ms = new MemoryStream();
            XmlDictionaryWriter writer = null;
            switch (messageFormat)
            {
                case WebContentFormat.Default:
                case WebContentFormat.Xml:
                    writer = XmlDictionaryWriter.CreateTextWriter(ms);
                    break;
                case WebContentFormat.Json:
                    writer = JsonReaderWriterFactory.CreateJsonWriter(ms);
                    break;
                case WebContentFormat.Raw:
                    // special case for raw, easier implemented separately 
                    return this.ReadRawBody(ref message);
            }

            message.WriteMessage(writer);
            writer.Flush();
            string messageBody = Encoding.UTF8.GetString(ms.ToArray());

            // Here would be a good place to change the message body, if so desired. 

            // now that the message was read, it needs to be recreated. 
            ms.Position = 0;

            // if the message body was modified, needs to reencode it, as show below 
            // ms = new MemoryStream(Encoding.UTF8.GetBytes(messageBody)); 

            XmlDictionaryReader reader;
            if (messageFormat == WebContentFormat.Json)
            {
                reader = JsonReaderWriterFactory.CreateJsonReader(ms, XmlDictionaryReaderQuotas.Max);
            }
            else
            {
                reader = XmlDictionaryReader.CreateTextReader(ms, XmlDictionaryReaderQuotas.Max);
            }

            Message newMessage = Message.CreateMessage(reader, int.MaxValue, message.Version);
            newMessage.Properties.CopyProperties(message.Properties);
            message = newMessage;

            return messageBody;
        }

        private string ReadRawBody(ref Message message)
        {
            XmlDictionaryReader bodyReader = message.GetReaderAtBodyContents();
            bodyReader.ReadStartElement("Binary");
            byte[] bodyBytes = bodyReader.ReadContentAsBase64();
            string messageBody = Encoding.UTF8.GetString(bodyBytes);

            // Now to recreate the message 
            MemoryStream ms = new MemoryStream();
            XmlDictionaryWriter writer = XmlDictionaryWriter.CreateBinaryWriter(ms);
            writer.WriteStartElement("Binary");
            writer.WriteBase64(bodyBytes, 0, bodyBytes.Length);
            writer.WriteEndElement();
            writer.Flush();
            ms.Position = 0;
            XmlDictionaryReader reader = XmlDictionaryReader.CreateBinaryReader(ms, XmlDictionaryReaderQuotas.Max);
            Message newMessage = Message.CreateMessage(reader, int.MaxValue, message.Version);
            newMessage.Properties.CopyProperties(message.Properties);
            message = newMessage;

            return messageBody;
        }
        #endregion

        #region AppDomain methods
        public static void SetCorrelationID(Guid correlationID)
        {
            AppDomain.CurrentDomain.SetData(HEADER_NAMESPACE + ":" + HEADER_CORRELATION, correlationID);                
        }
        public static Guid GetCorrelationID()
        {
            if (AppDomain.CurrentDomain.GetData(HEADER_NAMESPACE + ":" + HEADER_CORRELATION) != null)
                return (Guid) AppDomain.CurrentDomain.GetData(HEADER_NAMESPACE + ":" + HEADER_CORRELATION);
            
            return Guid.NewGuid();
        }
        private static void SetEventName(string name)
        {
            AppDomain.CurrentDomain.SetData(HEADER_NAMESPACE + ":" + HEADER_EVENT, name);                
        }
        public static string GetEventName()
        {
            if (AppDomain.CurrentDomain.GetData(HEADER_NAMESPACE + ":" + HEADER_EVENT) != null)
                return (string) AppDomain.CurrentDomain.GetData(HEADER_NAMESPACE + ":" + HEADER_EVENT);

            return string.Empty;
        }
        #endregion
    }
}